package com.studyinghome.config;

import com.studyinghome.utils.IpUtil;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.time.LocalDate;
import java.util.Objects;
import java.util.concurrent.TimeUnit;

/**
 * @author Leslie (panxiang_work@163.com)
 * @website https://www.studyinghome.com
 * @create 2023-04-09 20:15
 */
@Component
public class RedisLimit {
    @Autowired
    private RedisTemplate<String, Integer> redisTemplate;

    /**
     * 限制
     *
     * @param prefix     redis key的前缀
     * @param limitCount 限流次数
     * @param time       限流时间
     * @param ipLimit    开启IP限制
     * @param ipCount    ip限制数
     * @param ipTimeout  ip限制时间
     * @return true:通过 false:被限流
     */
    public boolean limit(String prefix, int limitCount, int time, boolean ipLimit, int ipCount, int ipTimeout) {
        boolean limitIp = true;
        if (ipLimit) {
            // 对IP进行限制
            ServletRequestAttributes attributes = ((ServletRequestAttributes) RequestContextHolder
                    .getRequestAttributes());
            HttpServletRequest request = null;
            if (attributes != null) {
                request = attributes.getRequest();
            }

            String ipAddress = IpUtil.getIpAddress(request);
            String key = ipAddress + ":" + LocalDate.now() + ":" + System.currentTimeMillis() / (1000 * 60 * ipTimeout);

            Long count = redisTemplate.opsForValue().increment(key, 1);
            redisTemplate.expire(key, time, TimeUnit.MINUTES);
            count = Objects.requireNonNullElse(count, 0L);
            limitIp = count <= ipCount;
        }

        String key = prefix + ":" + LocalDate.now() + ":" + System.currentTimeMillis() / (1000 * 60 * time);

        Long count = redisTemplate.opsForValue().increment(key, 1);
        redisTemplate.expire(key, time, TimeUnit.MINUTES);
        count = Objects.requireNonNullElse(count, 0L); // provide a default value of 0 if count is null
        return (count <= limitCount) && limitIp;
    }
}
